import datetime
import json
import logging
import os # Added import os
from collections import defaultdict
from collections.abc import Iterator, Sequence
from json import JSONDecodeError
from typing import Optional

from pydantic import BaseModel, ConfigDict, Field

from constants import HIDDEN_VALUE # Assuming constants is at the root of backend
from core.entities.model_entities import ModelStatus, ModelWithProviderEntity, SimpleModelProviderEntity
from core.entities.provider_entities import (
    CustomConfiguration,
    ModelSettings,
    SystemConfiguration,
    SystemConfigurationStatus,
)
from core.helper import encrypter
from core.helper.model_provider_cache import ProviderCredentialsCache, ProviderCredentialsCacheType
from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType
from core.model_runtime.entities.provider_entities import (
    ConfigurateMethod,
    CredentialFormSchema,
    FormType,
    ProviderEntity,
)
from core.model_runtime.model_providers.__base.ai_model import AIModel
from core.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
from core.plugin.entities.plugin import ModelProviderID
# from extensions.ext_database import db # Removed Dify specific db import
from extensions import db # Use the db instance from our extensions.py
from models.provider import (
    LoadBalancingModelConfig,
    Provider,
    ProviderModel,
    ProviderModelSetting,
    ProviderType,
    TenantPreferredModelProvider,
)

logger = logging.getLogger(__name__)

original_provider_configurate_methods: dict[str, list[ConfigurateMethod]] = {}

# Placeholder for tenant_id in a standalone context. This might need a more robust solution.
STANDALONE_TENANT_ID = os.getenv("STANDALONE_TENANT_ID", "standalone_tenant")

class ProviderConfiguration(BaseModel):
    """
    Model class for provider configuration.
    """

    tenant_id: str = Field(default=STANDALONE_TENANT_ID) # Default tenant_id for standalone
    provider: ProviderEntity
    preferred_provider_type: ProviderType
    using_provider_type: ProviderType
    system_configuration: SystemConfiguration
    custom_configuration: CustomConfiguration
    model_settings: list[ModelSettings]

    model_config = ConfigDict(protected_namespaces=())

    def __init__(self, **data):
        if "tenant_id" not in data:
            data["tenant_id"] = STANDALONE_TENANT_ID
        super().__init__(**data)

        if self.provider.provider not in original_provider_configurate_methods:
            original_provider_configurate_methods[self.provider.provider] = []
            for configurate_method in self.provider.configurate_methods:
                original_provider_configurate_methods[self.provider.provider].append(configurate_method)

        if original_provider_configurate_methods[self.provider.provider] == [ConfigurateMethod.CUSTOMIZABLE_MODEL]:
            if (
                any(
                    len(quota_configuration.restrict_models) > 0
                    for quota_configuration in self.system_configuration.quota_configurations
                )
                and ConfigurateMethod.PREDEFINED_MODEL not in self.provider.configurate_methods
            ):
                self.provider.configurate_methods.append(ConfigurateMethod.PREDEFINED_MODEL)

    def get_current_credentials(self, model_type: ModelType, model: str) -> Optional[dict]:
        if self.model_settings:
            for model_setting in self.model_settings:
                if model_setting.model_type == model_type and model_setting.model == model:
                    if not model_setting.enabled:
                        raise ValueError(f"Model {model} is disabled.")

        if self.using_provider_type == ProviderType.SYSTEM:
            restrict_models = []
            for quota_configuration in self.system_configuration.quota_configurations:
                if self.system_configuration.current_quota_type != quota_configuration.quota_type:
                    continue
                restrict_models = quota_configuration.restrict_models

            copy_credentials = (
                self.system_configuration.credentials.copy() if self.system_configuration.credentials else {}
            )
            if restrict_models:
                for restrict_model in restrict_models:
                    if (
                        restrict_model.model_type == model_type
                        and restrict_model.model == model
                        and restrict_model.base_model_name
                    ):
                        copy_credentials["base_model_name"] = restrict_model.base_model_name
            return copy_credentials
        else:
            credentials = None
            if self.custom_configuration.models:
                for model_configuration in self.custom_configuration.models:
                    if model_configuration.model_type == model_type and model_configuration.model == model:
                        credentials = model_configuration.credentials
                        break
            if not credentials and self.custom_configuration.provider:
                credentials = self.custom_configuration.provider.credentials
            return credentials

    def get_system_configuration_status(self) -> Optional[SystemConfigurationStatus]:
        if self.system_configuration.enabled is False:
            return SystemConfigurationStatus.UNSUPPORTED
        current_quota_type = self.system_configuration.current_quota_type
        current_quota_configuration = next(
            (q for q in self.system_configuration.quota_configurations if q.quota_type == current_quota_type), None
        )
        if current_quota_configuration is None:
            return None
        if not current_quota_configuration:
            return SystemConfigurationStatus.UNSUPPORTED
        return (
            SystemConfigurationStatus.ACTIVE
            if current_quota_configuration.is_valid
            else SystemConfigurationStatus.QUOTA_EXCEEDED
        )

    def is_custom_configuration_available(self) -> bool:
        return self.custom_configuration.provider is not None or len(self.custom_configuration.models) > 0

    def get_custom_credentials(self, obfuscated: bool = False) -> dict | None:
        if self.custom_configuration.provider is None:
            return None
        credentials = self.custom_configuration.provider.credentials
        if not obfuscated:
            return credentials
        return self.obfuscated_credentials(
            credentials=credentials,
            credential_form_schemas=self.provider.provider_credential_schema.credential_form_schemas
            if self.provider.provider_credential_schema
            else [],
        )

    def _get_custom_provider_credentials(self) -> Provider | None:
        model_provider_id = ModelProviderID(self.provider.provider)
        provider_names = [self.provider.provider]
        if model_provider_id.is_langgenius():
            provider_names.append(model_provider_id.provider_name)
        provider_record = (
            db.session.query(Provider)
            .filter(
                Provider.tenant_id == self.tenant_id,
                Provider.provider_type == ProviderType.CUSTOM.value,
                Provider.provider_name.in_(provider_names),
            )
            .first()
        )
        return provider_record

    def custom_credentials_validate(self, credentials: dict) -> tuple[Provider | None, dict]:
        provider_record = self._get_custom_provider_credentials()
        provider_credential_secret_variables = self.extract_secret_variables(
            self.provider.provider_credential_schema.credential_form_schemas
            if self.provider.provider_credential_schema
            else []
        )
        if provider_record:
            try:
                if provider_record.encrypted_config:
                    if not provider_record.encrypted_config.startswith("{"):
                        original_credentials = {"openai_api_key": provider_record.encrypted_config}
                    else:
                        original_credentials = json.loads(provider_record.encrypted_config)
                else:
                    original_credentials = {}
            except JSONDecodeError:
                original_credentials = {}
            for key, value in credentials.items():
                if key in provider_credential_secret_variables:
                    if value == HIDDEN_VALUE and key in original_credentials:
                        credentials[key] = encrypter.decrypt_token(self.tenant_id, original_credentials[key])
        model_provider_factory = ModelProviderFactory(self.tenant_id)
        credentials = model_provider_factory.provider_credentials_validate(
            provider=self.provider.provider, credentials=credentials
        )
        for key, value in credentials.items():
            if key in provider_credential_secret_variables:
                credentials[key] = encrypter.encrypt_token(self.tenant_id, value)
        return provider_record, credentials

    def add_or_update_custom_credentials(self, credentials: dict) -> None:
        provider_record, validated_credentials = self.custom_credentials_validate(credentials)
        if provider_record:
            provider_record.encrypted_config = json.dumps(validated_credentials)
            provider_record.is_valid = True
            provider_record.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
        else:
            provider_record = Provider(
                tenant_id=self.tenant_id,
                provider_name=self.provider.provider,
                provider_type=ProviderType.CUSTOM.value,
                encrypted_config=json.dumps(validated_credentials),
                is_valid=True
            )
            db.session.add(provider_record)
        db.session.commit()
        provider_model_credentials_cache = ProviderCredentialsCache(
            tenant_id=self.tenant_id, identity_id=provider_record.id, cache_type=ProviderCredentialsCacheType.PROVIDER
        )
        provider_model_credentials_cache.delete()
        self.switch_preferred_provider_type(ProviderType.CUSTOM)

    def delete_custom_credentials(self) -> None:
        provider_record = self._get_custom_provider_credentials()
        if provider_record:
            self.switch_preferred_provider_type(ProviderType.SYSTEM)
            db.session.delete(provider_record)
            db.session.commit()
            provider_model_credentials_cache = ProviderCredentialsCache(
                tenant_id=self.tenant_id,
                identity_id=provider_record.id,
                cache_type=ProviderCredentialsCacheType.PROVIDER,
            )
            provider_model_credentials_cache.delete()

    def get_custom_model_credentials(
        self, model_type: ModelType, model: str, obfuscated: bool = False
    ) -> Optional[dict]:
        if not self.custom_configuration.models:
            return None
        for model_configuration in self.custom_configuration.models:
            if model_configuration.model_type == model_type and model_configuration.model == model:
                credentials = model_configuration.credentials
                if not obfuscated:
                    return credentials
                return self.obfuscated_credentials(
                    credentials=credentials,
                    credential_form_schemas=self.provider.model_credential_schema.credential_form_schemas
                    if self.provider.model_credential_schema
                    else [],
                )
        return None

    def _get_custom_model_credentials(
        self,
        model_type: ModelType,
        model: str,
    ) -> ProviderModel | None:
        model_provider_id = ModelProviderID(self.provider.provider)
        provider_names = [self.provider.provider]
        if model_provider_id.is_langgenius():
            provider_names.append(model_provider_id.provider_name)
        provider_model_record = (
            db.session.query(ProviderModel)
            .filter(
                ProviderModel.tenant_id == self.tenant_id,
                ProviderModel.provider_name.in_(provider_names),
                ProviderModel.model_name == model,
                ProviderModel.model_type == model_type.to_origin_model_type(),
            )
            .first()
        )
        return provider_model_record

    def custom_model_credentials_validate(
        self, model_type: ModelType, model: str, credentials: dict
    ) -> tuple[ProviderModel | None, dict]:
        provider_model_record = self._get_custom_model_credentials(model_type, model)
        provider_credential_secret_variables = self.extract_secret_variables(
            self.provider.model_credential_schema.credential_form_schemas
            if self.provider.model_credential_schema
            else []
        )
        if provider_model_record:
            try:
                original_credentials = (
                    json.loads(provider_model_record.encrypted_config) if provider_model_record.encrypted_config else {}
                )
            except JSONDecodeError:
                original_credentials = {}
            for key, value in credentials.items():
                if key in provider_credential_secret_variables:
                    if value == HIDDEN_VALUE and key in original_credentials:
                        credentials[key] = encrypter.decrypt_token(self.tenant_id, original_credentials[key])
        model_provider_factory = ModelProviderFactory(self.tenant_id)
        credentials = model_provider_factory.model_credentials_validate(
            provider=self.provider.provider, model_type=model_type, model=model, credentials=credentials
        )
        for key, value in credentials.items():
            if key in provider_credential_secret_variables:
                credentials[key] = encrypter.encrypt_token(self.tenant_id, value)
        return provider_model_record, credentials

    def add_or_update_custom_model_credentials(self, model_type: ModelType, model: str, credentials: dict) -> None:
        provider_model_record, validated_credentials = self.custom_model_credentials_validate(model_type, model, credentials)
        if provider_model_record:
            provider_model_record.encrypted_config = json.dumps(validated_credentials)
            provider_model_record.is_valid = True
            provider_model_record.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
        else:
            provider_model_record = ProviderModel(
                tenant_id=self.tenant_id,
                provider_name=self.provider.provider,
                model_name=model,
                model_type=model_type.to_origin_model_type(),
                encrypted_config=json.dumps(validated_credentials),
                is_valid=True
            )
            db.session.add(provider_model_record)
        db.session.commit()
        provider_model_credentials_cache = ProviderCredentialsCache(
            tenant_id=self.tenant_id,
            identity_id=f"{provider_model_record.provider_name}:{provider_model_record.model_name}:{provider_model_record.model_type}",
            cache_type=ProviderCredentialsCacheType.MODEL,
        )
        provider_model_credentials_cache.delete()

    def delete_custom_model_credentials(self, model_type: ModelType, model: str) -> None:
        provider_model_record = self._get_custom_model_credentials(model_type, model)
        if provider_model_record:
            db.session.delete(provider_model_record)
            db.session.commit()
            provider_model_credentials_cache = ProviderCredentialsCache(
                tenant_id=self.tenant_id,
                identity_id=f"{provider_model_record.provider_name}:{provider_model_record.model_name}:{provider_model_record.model_type}",
                cache_type=ProviderCredentialsCacheType.MODEL,
            )
            provider_model_credentials_cache.delete()

    def switch_preferred_provider_type(self, provider_type: ProviderType) -> None:
        tenant_preferred_model_provider = (
            db.session.query(TenantPreferredModelProvider)
            .filter(
                TenantPreferredModelProvider.tenant_id == self.tenant_id,
                TenantPreferredModelProvider.provider_name == self.provider.provider,
            )
            .first()
        )
        if tenant_preferred_model_provider:
            tenant_preferred_model_provider.preferred_provider_type = provider_type.value
            tenant_preferred_model_provider.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
        else:
            tenant_preferred_model_provider = TenantPreferredModelProvider(
                tenant_id=self.tenant_id,
                provider_name=self.provider.provider,
                preferred_provider_type=provider_type.value,
            )
            db.session.add(tenant_preferred_model_provider)
        db.session.commit()

    @classmethod
    def obfuscated_credentials(cls, credentials: dict, credential_form_schemas: list[CredentialFormSchema]) -> dict:
        obfuscated_credentials = credentials.copy()
        for credential_form_schema in credential_form_schemas:
            if credential_form_schema.type == FormType.SECRET_INPUT:
                if credential_form_schema.variable in obfuscated_credentials:
                    obfuscated_credentials[credential_form_schema.variable] = HIDDEN_VALUE
        return obfuscated_credentials

    @classmethod
    def extract_secret_variables(cls, credential_form_schemas: list[CredentialFormSchema]) -> list[str]:
        secret_variables = []
        for credential_form_schema in credential_form_schemas:
            if credential_form_schema.type == FormType.SECRET_INPUT:
                secret_variables.append(credential_form_schema.variable)
        return secret_variables

class ProviderModelBundle(BaseModel):
    provider: str
    model_type: ModelType
    models: list[ModelWithProviderEntity]

    # pydantic configs
    model_config = ConfigDict(protected_namespaces=())

    @classmethod
    def from_provider_entities(
        cls,
        tenant_id: str,
        provider_entities: Sequence[ProviderEntity],
        model_type_filter: Optional[ModelType] = None,
        fetch_from: FetchFrom = FetchFrom.CONFIG,
    ) -> Iterator["ProviderModelBundle"]:
        model_provider_factory = ModelProviderFactory(tenant_id)
        for provider_entity in provider_entities:
            if not provider_entity.supported_model_types:
                continue

            if model_type_filter and model_type_filter not in provider_entity.supported_model_types:
                continue

            for model_type in provider_entity.supported_model_types:
                if model_type_filter and model_type != model_type_filter:
                    continue

                provider_configuration = model_provider_factory.get_provider_configuration(
                    provider=provider_entity.provider, fetch_from=fetch_from
                )

                if not provider_configuration:
                    continue

                models: list[ModelWithProviderEntity] = []
                for model_entity in model_provider_factory.list_models(
                    provider=provider_entity.provider, model_type=model_type, fetch_from=fetch_from
                ):
                    model_status = ModelStatus.ACTIVE
                    if provider_configuration.system_configuration.enabled is False:
                        model_status = ModelStatus.NO_CONFIGURE
                    elif provider_configuration.using_provider_type == ProviderType.SYSTEM:
                        system_configuration_status = provider_configuration.get_system_configuration_status()
                        if system_configuration_status == SystemConfigurationStatus.QUOTA_EXCEEDED:
                            model_status = ModelStatus.QUOTA_EXCEEDED
                        elif system_configuration_status == SystemConfigurationStatus.UNSUPPORTED:
                            model_status = ModelStatus.NO_CONFIGURE
                    elif provider_configuration.using_provider_type == ProviderType.CUSTOM:
                        if not provider_configuration.is_custom_configuration_available():
                            model_status = ModelStatus.NO_CONFIGURE

                    # check if model is disabled by admin
                    if provider_configuration.model_settings:
                        for model_setting in provider_configuration.model_settings:
                            if (
                                model_setting.model_type == model_type
                                and model_setting.model == model_entity.model
                                and not model_setting.enabled
                            ):
                                model_status = ModelStatus.NO_CONFIGURE
                                break

                    models.append(
                        ModelWithProviderEntity(
                            provider=SimpleModelProviderEntity(
                                provider=provider_entity.provider,
                                label=provider_entity.label,
                                icon_small=provider_entity.icon_small,
                                icon_large=provider_entity.icon_large,
                            ),
                            model=model_entity.model,
                            label=model_entity.label,
                            model_type=model_type,
                            status=model_status,
                            fetch_from=model_entity.fetch_from,
                            model_properties=model_entity.model_properties,
                            deprecated=model_entity.deprecated,
                        )
                    )

                if models:
                    yield ProviderModelBundle(provider=provider_entity.provider, model_type=model_type, models=models)

    @classmethod
    def get_configured_provider_model_bundle(
        cls, tenant_id: str, provider: str, model_type: ModelType, fetch_from: FetchFrom = FetchFrom.CONFIG
    ) -> Optional["ProviderModelBundle"]:
        model_provider_factory = ModelProviderFactory(tenant_id)
        provider_entity = model_provider_factory.get_provider(provider=provider, fetch_from=fetch_from)
        if not provider_entity:
            return None

        if model_type not in provider_entity.supported_model_types:
            return None

        provider_configuration = model_provider_factory.get_provider_configuration(
            provider=provider_entity.provider, fetch_from=fetch_from
        )

        if not provider_configuration:
            return None

        models: list[ModelWithProviderEntity] = []
        for model_entity in model_provider_factory.list_models(
            provider=provider_entity.provider, model_type=model_type, fetch_from=fetch_from
        ):
            model_status = ModelStatus.ACTIVE
            if provider_configuration.system_configuration.enabled is False:
                model_status = ModelStatus.NO_CONFIGURE
            elif provider_configuration.using_provider_type == ProviderType.SYSTEM:
                system_configuration_status = provider_configuration.get_system_configuration_status()
                if system_configuration_status == SystemConfigurationStatus.QUOTA_EXCEEDED:
                    model_status = ModelStatus.QUOTA_EXCEEDED
                elif system_configuration_status == SystemConfigurationStatus.UNSUPPORTED:
                    model_status = ModelStatus.NO_CONFIGURE
            elif provider_configuration.using_provider_type == ProviderType.CUSTOM:
                if not provider_configuration.is_custom_configuration_available():
                    model_status = ModelStatus.NO_CONFIGURE

            # check if model is disabled by admin
            if provider_configuration.model_settings:
                for model_setting in provider_configuration.model_settings:
                    if (
                        model_setting.model_type == model_type
                        and model_setting.model == model_entity.model
                        and not model_setting.enabled
                    ):
                        model_status = ModelStatus.NO_CONFIGURE
                        break

            models.append(
                ModelWithProviderEntity(
                    provider=SimpleModelProviderEntity(
                        provider=provider_entity.provider,
                        label=provider_entity.label,
                        icon_small=provider_entity.icon_small,
                        icon_large=provider_entity.icon_large,
                    ),
                    model=model_entity.model,
                    label=model_entity.label,
                    model_type=model_type,
                    status=model_status,
                    fetch_from=model_entity.fetch_from,
                    model_properties=model_entity.model_properties,
                    deprecated=model_entity.deprecated,
                )
            )

        if models:
            return ProviderModelBundle(provider=provider_entity.provider, model_type=model_type, models=models)

        return None

    @classmethod
    def get_provider_model(cls, tenant_id: str, provider: str, model: str) -> Optional[AIModelEntity]:
        """
        Get provider model.

        :param tenant_id: tenant id
        :param provider: provider name
        :param model: model name
        :return:
        """
        model_provider_factory = ModelProviderFactory(tenant_id)
        provider_entity = model_provider_factory.get_provider(provider=provider)
        if not provider_entity:
            return None

        for model_type in provider_entity.supported_model_types:
            for model_entity in model_provider_factory.list_models(provider=provider, model_type=model_type):
                if model_entity.model == model:
                    return model_entity

        return None

    @classmethod
    def get_provider_model_valid_labels(cls, tenant_id: str) -> dict[str, dict[str, list[str]]]:
        """
        Get provider model valid labels.

        :param tenant_id: tenant id
        :return:
        """
        provider_model_valid_labels = defaultdict(lambda: defaultdict(list))
        for provider_model_bundle in cls.from_provider_entities(tenant_id=tenant_id):
            for model_with_provider_entity in provider_model_bundle.models:
                if model_with_provider_entity.status == ModelStatus.ACTIVE:
                    provider_model_valid_labels[provider_model_bundle.provider][
                        provider_model_bundle.model_type.value
                    ].append(model_with_provider_entity.model)

        return provider_model_valid_labels




class ProviderConfigurations(BaseModel):
    """
    Model class for a collection of provider configurations.
    """
    tenant_id: str = Field(default=STANDALONE_TENANT_ID)
    providers: dict[str, ProviderConfiguration] = Field(default_factory=dict)

    model_config = ConfigDict(protected_namespaces=())

    def __init__(self, tenant_id: str, providers: Optional[dict[str, ProviderConfiguration]] = None):
        super().__init__(tenant_id=tenant_id, providers=providers if providers else {})

    def __getitem__(self, key: str) -> ProviderConfiguration:
        return self.providers[key]

    def __setitem__(self, key: str, value: ProviderConfiguration) -> None:
        self.providers[key] = value

    def __contains__(self, key: str) -> bool:
        return key in self.providers

    def get(self, key: str, default: Optional[ProviderConfiguration] = None) -> Optional[ProviderConfiguration]:
        return self.providers.get(key, default)

    def get_models(
        self, model_type: Optional[ModelType] = None, only_active: bool = False
    ) -> list[ModelWithProviderEntity]:
        # Placeholder implementation, needs to be fully fleshed out based on Dify's logic
        all_models: list[ModelWithProviderEntity] = []
        for provider_config in self.providers.values():
            # This is a simplified placeholder. The actual logic to get models from a 
            # ProviderConfiguration instance would be more complex and involve its methods.
            # For now, we'll assume a ProviderConfiguration might have a method or property 
            # to list its models, or we iterate through its model_settings if applicable.
            # This part needs careful refactoring based on Dify's original structure.
            if provider_config.provider and provider_config.provider.models:
                for model_entity in provider_config.provider.models: # Assuming ProviderEntity has a models list
                    if model_type and model_entity.model_type != model_type:
                        continue
                    # Add status check if only_active is True
                    # This is a very basic placeholder for model listing
                    all_models.append(
                        ModelWithProviderEntity(
                            provider=SimpleModelProviderEntity(
                                provider=provider_config.provider.provider,
                                label=provider_config.provider.label,
                                icon_small=provider_config.provider.icon_small,
                                icon_large=provider_config.provider.icon_large,
                            ),
                            model=model_entity.model, # model name
                            label=model_entity.label, # model label
                            model_type=model_entity.model_type, # model type
                            status=ModelStatus.ACTIVE, # Placeholder status
                            fetch_from=FetchFrom.CONFIG, # Placeholder fetch_from
                            model_properties=model_entity.model_properties if hasattr(model_entity, 'model_properties') else {},
                            deprecated=False # Placeholder deprecated
                        )
                    )
        return all_models
